""" 负责socket连接与心跳 """

import socket
import json

from adb import actions

divider = b'\x18'
heart = b'\x01'
end = b'\x00'


class AdbConnection(object):
    """ 表示一个连接 """
    def __init__(self, address='localhost', port=9318):
        self.socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        self.port = port
        self.address = address

    def connect(self):
        self.socket.connect((self.address, self.port))
        self.socket.settimeout(20)

    def send(self, action_obj: actions.ActionObj):
        if not isinstance(action_obj, actions.ActionObj):
            raise RuntimeError("action_obj must be a instance of adb.actions.ActionObj")
        return self._send(self.socket, action_obj.ACTION, action_obj, action_obj.RETURN_BYTES)

    @staticmethod
    def _send(socket, action, content, res_bytes):
        if isinstance(content, bytes):
            result_content = content
        elif isinstance(content, str):
            result_content = bytes(content, 'utf8')
        elif isinstance(content, dict):
            result_content = bytes(json.dumps(content, ensure_ascii=False), 'utf8')
        else:
            result_content = bytes(json.dumps(content.__dict__, ensure_ascii=False), 'utf8')

        length = len(result_content) if result_content is not None else 0
        socket.sendall(divider)
        socket.sendall(action)
        socket.sendall(length.to_bytes(2, 'big', signed=False))
        if result_content is not None and len(result_content) != 0:
            socket.sendall(result_content)
        while True:
            next_divider = socket.recv(1)
            if next_divider == "":
                return None
            if next_divider == heart:
                socket.sendall(heart)
                print("send heart:")
                continue
            if next_divider == divider:
                break
        response_length = int.from_bytes(socket.recv(2), 'big', signed=False)
        response = None
        if response_length != 0:
            response_bytes = socket.recv(response_length)
            if res_bytes:
                return response_bytes
            response = json.loads(str(response_bytes, encoding="utf8"))
        return response

    def close(self):
        self.socket.close()
